"""
模型下载脚本
用于下载预训练的车牌检测和识别模型
"""

import sys
import os
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.models import ModelLoader
from src.utils import Logger, Config


def main():
    """
    主函数
    """
    # 初始化日志
    logger = Logger.get_logger()
    logger.info("开始下载模型...")
    
    # 加载配置
    config = Config(str(project_root / "config.yaml"))
    
    # 创建模型加载器
    loader = ModelLoader(
        models_dir=str(project_root / "models"),
        config=config.get_all()
    )
    
    # 列出所有可用模型
    models = loader.list_models()
    logger.info("可用模型列表:")
    for name, info in models.items():
        status = "已安装" if info['installed'] else "未安装"
        logger.info(f"  - {name}: {info['description']} ({info['size']}) [{status}]")
    
    # 下载未安装的模型
    for name, info in models.items():
        if not info['installed']:
            logger.info(f"正在下载 {name}...")
            model_path = loader.load_model(name)
            if model_path:
                logger.success(f"{name} 下载成功: {model_path}")
            else:
                logger.error(f"{name} 下载失败")
    
    logger.info("模型下载完成")
    
    # 再次列出模型状态
    models = loader.list_models()
    logger.info("\n最终模型状态:")
    for name, info in models.items():
        status = "已安装" if info['installed'] else "未安装"
        logger.info(f"  - {name}: [{status}]")


if __name__ == "__main__":
    main() 